package demo.hdz.mediacodecdecode;

import demo.hdz.mediacodecdecode.bean.VideoSizeBean;

/**
 * @author hdz
 * @date 2019/2/19
 *
 * 参考：
 *      https://github.com/leixiaohua1020/simplest_librtmp_example/blob/master/simplest_librtmp_send264/sps_decode.h
 *
 *      https://www.jianshu.com/p/f4f3b0764cec
 *      https://www.jianshu.com/p/06e7c742de44
 *      https://www.jianshu.com/p/26326ea02e69
 */

public class SpsDecode {

    private static final int START_CODE_LEN = 4*8;
    private static int mStartBit = 0;

    public SpsDecode() {
        mStartBit = START_CODE_LEN;
    }

    public VideoSizeBean decode(byte[] sps) {
        int width = 0;
        int height = 0;

        int forbidden_zero_bit = u(sps, 1, mStartBit);
        mStartBit += 1;
        int nal_ref_idc = u(sps, 2, mStartBit);
        mStartBit += 2;
        int nal_unit_type = u(sps, 5, mStartBit);
        mStartBit += 5;
        if (nal_unit_type == 7) {
            int profile_idc=u(sps,8, mStartBit);
            mStartBit += 8;
            int constraint_set0_flag=u(sps,1,mStartBit);//(buf[1] & 0x80)>>7;
            mStartBit += 1;
            int constraint_set1_flag=u(sps, 1, mStartBit);//(buf[1] & 0x40)>>6;
            mStartBit += 1;
            int constraint_set2_flag=u(sps, 1, mStartBit);//(buf[1] & 0x20)>>5;
            mStartBit += 1;
            int reserved_zero_5bits=u(sps, 5, mStartBit);
            mStartBit += 5;
            int level_idc=u(sps,8,mStartBit);//56
            mStartBit += 8;
            int seq_parameter_set_id = ue(sps, mStartBit);

            if( profile_idc == 100 || profile_idc == 110 ||
                    profile_idc == 122 || profile_idc == 144 ) {
                int chroma_format_idc=ue(sps,-1);
                if( chroma_format_idc == 3 ) {
                    int residual_colour_transform_flag = u(sps, 1, mStartBit);
                    mStartBit += 1;
                }
                int bit_depth_luma_minus8=ue(sps, -1);
                int bit_depth_chroma_minus8=ue(sps, -1);
                int qpprime_y_zero_transform_bypass_flag=u(sps, 1, mStartBit);
                mStartBit += 1;
                int seq_scaling_matrix_present_flag=u(sps, 1, mStartBit);
                mStartBit += 1;

                int[] seq_scaling_list_present_flag = new int[8];
                if( seq_scaling_matrix_present_flag != 0) {
                    for( int i = 0; i < 8; i++ ) {
                        seq_scaling_list_present_flag[i]=u(sps, 1, mStartBit);
                        mStartBit += 1;
                    }
                }
            }
            int log2_max_frame_num_minus4=ue(sps, mStartBit);
            int pic_order_cnt_type=ue(sps, mStartBit);
            if( pic_order_cnt_type == 0 ) {
                int log2_max_pic_order_cnt_lsb_minus4 = ue(sps, mStartBit);
            } else if( pic_order_cnt_type == 1 ) {
                int delta_pic_order_always_zero_flag=u(sps, 1, mStartBit);
                mStartBit += 1;
                int offset_for_non_ref_pic=se(sps, mStartBit);
                int offset_for_top_to_bottom_field=se(sps, mStartBit);
                int num_ref_frames_in_pic_order_cnt_cycle=ue(sps, mStartBit);

                int[] offset_for_ref_frame = new int[num_ref_frames_in_pic_order_cnt_cycle];
                for( int i = 0; i < num_ref_frames_in_pic_order_cnt_cycle; i++ ) {
                    offset_for_ref_frame[i] = se(sps, mStartBit);
                }
            }

            int num_ref_frames=ue(sps, mStartBit);
            int gaps_in_frame_num_value_allowed_flag=u(sps, 1, mStartBit);
            mStartBit += 1;
            int pic_width_in_mbs_minus1=ue(sps, mStartBit);
            int pic_height_in_map_units_minus1=ue(sps, mStartBit);

            width = (pic_width_in_mbs_minus1+1)*16;
            height = (pic_height_in_map_units_minus1+1)*16;
            System.out.println("width = "+width);
            System.out.println("height = "+height);

            int frame_mbs_only_flag=u(sps, 1, mStartBit);
            mStartBit += 1;
            if(frame_mbs_only_flag == 0) {
                int mb_adaptive_frame_field_flag = u(sps, 1, mStartBit);
                mStartBit += 1;
            }

            int direct_8x8_inference_flag=u(sps, 1, mStartBit);
            mStartBit += 1;
            int frame_cropping_flag=u(sps, 1, mStartBit);
            mStartBit += 1;
            if(frame_cropping_flag != 0) {
                int frame_crop_left_offset=ue(sps, mStartBit);
                int frame_crop_right_offset=ue(sps, mStartBit);
                int frame_crop_top_offset=ue(sps, mStartBit);
                int frame_crop_bottom_offset=ue(sps, mStartBit);
            }
            int vui_parameter_present_flag=u(sps, 1, mStartBit);
            mStartBit += 1;

            if(vui_parameter_present_flag != 0) {
                int aspect_ratio_info_present_flag=u(sps, 1, mStartBit);
                mStartBit += 1;
                if(aspect_ratio_info_present_flag != 0) {
                    int aspect_ratio_idc=u(sps, 8, mStartBit);
                    mStartBit += 8;
                    if(aspect_ratio_idc==255) {
                        int sar_width=u(sps, 16, mStartBit);
                        mStartBit += 16;
                        int sar_height=u(sps, 16, mStartBit);
                        mStartBit += 16;
                    }
                }
                int overscan_info_present_flag=u(sps, 1, mStartBit);
                mStartBit += 1;
                if(overscan_info_present_flag != 0) {
                    int overscan_appropriate_flagu = u(sps, 1, mStartBit);
                    mStartBit += 1;
                }
                int video_signal_type_present_flag=u(sps, 1, mStartBit);
                mStartBit += 1;
                if(video_signal_type_present_flag != 0) {
                    int video_format=u(sps, 3, mStartBit);
                    mStartBit += 3;
                    int video_full_range_flag=u(sps, 1, mStartBit);
                    mStartBit += 1;
                    int colour_description_present_flag=u(sps, 1, mStartBit);
                    mStartBit += 1;
                    if(colour_description_present_flag != 0) {
                        int colour_primaries=u(sps, 8, mStartBit);
                        mStartBit += 8;
                        int transfer_characteristics=u(sps, 8, mStartBit);
                        mStartBit += 8;
                        int matrix_coefficients=u(sps, 8, mStartBit);
                        mStartBit += 8;
                    }
                }
                int chroma_loc_info_present_flag=u(sps, 1, mStartBit);
                mStartBit += 1;
                if(chroma_loc_info_present_flag != 0) {
                    int chroma_sample_loc_type_top_field=ue(sps,mStartBit);
                    int chroma_sample_loc_type_bottom_field=ue(sps,mStartBit);
                }
                int timing_info_present_flag=u(sps, 1, mStartBit);
                mStartBit += 1;
                if(timing_info_present_flag != 0) {
                    int num_units_in_tick=u(sps, 32, mStartBit);
                    mStartBit += 32;
                    int time_scale=u(sps, 32, mStartBit);
                    mStartBit += 32;
                    int fps=time_scale/(2*num_units_in_tick);
                    System.out.println("fps = " + fps);
                }
                int nal_hrd_parameters_present_flag=u(sps, 1, mStartBit);
                mStartBit += 1;
                int ncl_hrd_parameters_present_flag=u(sps, 1, mStartBit);
                mStartBit += 1;
            }
        }
        return new VideoSizeBean(width, height);
    }



    /*
     * 从数据流data中第StartBit位开始读，读bitCnt位，以无符号整形返回
     */
    public static short u(byte[] data, int bitCnt, int StartBit){
        short ret = 0;
        int start = StartBit;
        for(int i = 0; i < bitCnt; i++){
            ret<<=1;
            if ((data[start / 8] & (0x80 >> (start%8))) != 0) {
                ret += 1;
            }
            start++;
        }
        return ret;
    }

    /*
     * 无符号指数哥伦布熵编码
     * leadingZeroBits = −1;
     * for( b = 0; !b; leadingZeroBits++ )
     *    b = read_bits( 1 )
     *
     * 变量codeNum 按照如下方式赋值：
     * codeNum = 2^leadingZeroBits − 1 + read_bits( leadingZeroBits )
     * 这里read_bits( leadingZeroBits )的返回值使用高位在先的二进制无符号整数表示。
     * read_bits( leadingZeroBits )为从中间1开始，往后顺序数leadingZeroBits个比特位。
     * 其中leadingZeroBits，就是码字中，第一个比特值为1的前面0的个数，也即前导零个数。
     * 而 read_bits( leadingZeroBits )，则是从这第一个1后面，再顺序读取leadingZeroBits个比特。
     */
    public short ue(byte[] data,int StartBit){
        short ret = 0;
        int leadingZeroBits = -1;

        //如果传入-1，那么就用上次记录的静态变量
        int tempStartBit = (StartBit == -1)?mStartBit:StartBit;

        for( int b = 0; b != 1; leadingZeroBits++ ) {//读到第一个不为0的数，计算前面0的个数
            b = u(data,1,tempStartBit++);
        }
        ret = (short) (Math.pow(2, leadingZeroBits) - 1 + u(data, leadingZeroBits, tempStartBit));
        mStartBit = tempStartBit + leadingZeroBits;
        return ret;
    }
    /*
     * 有符号指数哥伦布熵编码
     *
     * 9.1.1 有符号指数哥伦布编码的映射过程
     *按照9.1节规定，本过程的输入是codeNum。
     *本过程的输出是se(v)的值。
     *表9-3中给出了分配给codeNum的语法元素值的规则，语法元素值按照绝对值的升序排列，负值按照其绝对
     *值参与排列，但列在绝对值相等的正值之后。
     *表 9-3－有符号指数哥伦布编码语法元素se(v)值与codeNum的对应
     *codeNum 语法元素值
     *	0 		0
     *	1		1
     *	2		−1
     *	3		2
     *	4		−2
     *	5		3
     *	6		−3
     *	k		(−1)^(k+1) Ceil( k÷2 )
     */
    public int se(byte[] data,int StartBit){
        int ret = 0;
        short codeNum = ue(data,StartBit);
        ret = (int) (Math.pow(-1, codeNum + 1)*Math.ceil(codeNum/2));
        return ret;
    }
}
